import json
import os


def coco_json_combine(dataset_list, dataset_target):
    labels_path = os.path.join(dataset_target, 'annotations')
    if not os.path.exists(labels_path):
        os.makedirs(labels_path)
    train_img_point = 0
    train_anno_point = 0
    test_img_point = 0
    test_anno_point = 0
    val_img_point = 0
    val_anno_point = 0
    train_all = {'images': [], 'categories': [], 'annotations': []}
    test_all = {'images': [], 'categories': [], 'annotations': []}
    val_all = {'images': [], 'categories': [], 'annotations': []}
    for num in range(len(dataset_list)):
        # print('train_img_point:', train_img_point)
        # print('train_anno_point:', train_anno_point)
        # Process train set
        json_train = os.path.join(dataset_list[num], 'annotations', 'voc07_train.json')
        # Parse json content
        with open(json_train, 'r', encoding='utf-8') as train_01:
            json_train_content = json.load(train_01)
        train_img_num = json_train_content['images'][-1]['id']
        train_anno_num = json_train_content['annotations'][-1]['id']
        # Update image id
        for img_info in json_train_content['images']:
            img_info['id'] = train_img_point + img_info['id']
        for anno_info in json_train_content['annotations']:
            anno_info['image_id'] = train_img_point + anno_info['image_id']
            # print(anno_info['id'])
            anno_info['id'] = train_anno_point + anno_info['id']
        # Set up image_id
        train_img_point = train_img_point + train_img_num
        print(json_train_content['images'][-1]['id'])
        print('update train_img_point:', train_img_point)
        train_anno_point = train_anno_point + train_anno_num
        print(json_train_content['annotations'][-1]['id'])
        print('update train_anno_point:', train_anno_point)
        for key in train_all:
            train_all[key] = train_all[key] + json_train_content[key]
        # Process test set
        json_test = os.path.join(dataset_list[num], 'annotations', 'voc07_test.json')
        # Parse json content
        with open(json_test, 'r', encoding='utf-8') as test_01:
            json_test_content = json.load(test_01)
        test_img_num = json_test_content['images'][-1]['id']
        test_anno_num = json_test_content['annotations'][-1]['id']
        # Update image id
        for img_info in json_test_content['images']:
            img_info['id'] = test_img_point + img_info['id']
        for anno_info in json_test_content['annotations']:
            anno_info['image_id'] = test_img_point + anno_info['image_id']
            anno_info['id'] = test_anno_point + anno_info['id']
        # Set up image_id
        test_img_point = test_img_point + test_img_num
        # print(json_test_content['images'][-1]['id'])
        # print('update test_img_point:', test_img_point)
        test_anno_point = test_anno_point + test_anno_num
        # print(json_test_content['annotations'][-1]['id'])
        # print('update test_anno_point:', test_anno_point)
        for key in test_all:
            test_all[key] = test_all[key] + json_test_content[key]
        # Process val set
        json_val = os.path.join(dataset_list[num], 'annotations', 'voc07_val.json')
        # Parse json content
        with open(json_val, 'r', encoding='utf-8') as val_01:
            json_val_content = json.load(val_01)
        val_img_num = json_val_content['images'][-1]['id']
        val_anno_num = json_val_content['annotations'][-1]['id']
        # Update image id
        for img_info in json_val_content['images']:
            img_info['id'] = val_img_point + img_info['id']
        for anno_info in json_val_content['annotations']:
            anno_info['image_id'] = val_img_point + anno_info['image_id']
            anno_info['id'] = val_anno_point + anno_info['id']
        # Set up image_id
        val_img_point = val_img_point + val_img_num
        # print(json_val_content['images'][-1]['id'])
        # print('update val_img_point:', val_img_point)
        val_anno_point = val_anno_point + val_anno_num
        # print(json_val_content['annotations'][-1]['id'])
        # print('update val_anno_point:', val_anno_point)
        for key in val_all:
            val_all[key] = val_all[key] + json_val_content[key]
    # Output content in sum json
    json_train_all = os.path.join(dataset_target, 'annotations', 'voc07_train.json')
    with open(json_train_all, 'w', encoding='utf-8') as content_all:
        json.dump(train_all, content_all)
    json_test_all = os.path.join(dataset_target, 'annotations', 'voc07_test.json')
    with open(json_test_all, 'w', encoding='utf-8') as content_all:
        json.dump(test_all, content_all)
    json_val_all = os.path.join(dataset_target, 'annotations', 'voc07_val.json')
    with open(json_val_all, 'w', encoding='utf-8') as content_all:
        json.dump(val_all, content_all)


if __name__ == "__main__":
    dataset_1 = r'D:\Project\dataset\COCO\DataSet_01'
    dataset_2 = r'D:\Project\dataset\COCO\DataSet_02'
    dataset_3 = r'D:\Project\dataset\COCO\DataSet_03'
    dataset_4 = r'D:\Project\dataset\COCO\DataSet_04'
    dataset_list = [dataset_1, dataset_2, dataset_3]
    coco_json_combine(dataset_list, dataset_4)
    json_train = r'\voc07_train.json'
    json_test = r'\voc07_test.json'
    json_val = r'\voc07_val.json'
    # with open(r'D:\Project\dataset\for_test\fire\annotations' + json_train, 'r', encoding='utf-8') as jr:
    with open(r'D:\Project\dataset\COCO\DataSet_04\annotations' + json_train, 'r', encoding='utf-8') as jr:
        json_content = json.load(jr)
    # print(json_content.keys())
    # print(json_content['images'][0].keys())
    # print(json_content['categories'][0].keys())
    # print(json_content['annotations'][0].keys())
    print(json_content['images'][0]['id'])
    print(json_content['annotations'][0]['id'])
    # for img in json_content['images']:
    #     print(img)
    # for anno in json_content['annotations']:
    #     if anno['image_id'] == 274 or anno['image_id'] == 382:
    #         print(anno)
    # with open(r'D:\Project\dataset\COCO\DataSet_01\annotations\voc07_train.json', 'r', encoding='utf-8') as jr:
    #     json_content = json.load(jr)
    # print('DataSet_01:', json_content['images'][-1]['id'])
    # print('DataSet_01:', json_content['annotations'][-1]['id'])
    # with open(r'D:\Project\dataset\COCO\DataSet_02\annotations\voc07_train.json', 'r', encoding='utf-8') as jr:
    #     json_content = json.load(jr)
    # print('DataSet_02:', json_content['images'][-1]['id'])
    # print('DataSet_02:', json_content['annotations'][-1]['id'])
    # with open(r'D:\Project\dataset\COCO\DataSet_03\annotations\voc07_train.json', 'r', encoding='utf-8') as jr:
    #     json_content = json.load(jr)
    # print('DataSet_03:', json_content['images'][-1]['id'])
    # print('DataSet_03:', json_content['annotations'][-1]['id'])
    # with open(r'D:\Project\dataset\COCO\DataSet_03\annotations\voc07_train.json', 'r', encoding='utf-8') as jr:
    #     json_content = json.load(jr)
    #     print(json_content.keys())
    # print(json_content['images'][-1]['id'])
    # print(json_content['images'][0].keys())
    # print(json_content['categories'][0].keys())
    # print(json_content['annotations'][0].keys())
    # print(json_content['annotations'][-1]['id'])
    # for anno in json_content['annotations']:
    #     print(anno['id'])
    # json_annotations_02 = json_content['annotations'][0].keys()
    # print(json_annotations_01 == json_annotations_02)
